import math
from typing import List

from langchain_core.embeddings import Embeddings
from zhipuai import ZhipuAI
from library.common.constants import ZHIPUAI_API_KEY


class ZhipuEmbeddings(Embeddings):
    def __init__(self):
        self.client = ZhipuAI(api_key=ZHIPUAI_API_KEY)  # 请填写您自己的APIKey
        self.batch_size = 10

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        result = []
        # 一个批次太多数据会报错
        for i in range(math.ceil(len(texts) / self.batch_size)):
            batch = texts[i * self.batch_size:(i + 1) * self.batch_size]
            response = self.client.embeddings.create(
                model="embedding-2",  # 填写需要调用的模型名称
                input=batch,
            )
            result.extend([item.embedding for item in response.data])
        return result

    def embed_query(self, text: str) -> List[float]:
        print(f"embed_query: {text}")
        response = self.client.embeddings.create(
            model="embedding-2",  # 填写需要调用的模型名称
            input=text,
        )
        return response.data[0].embedding


if __name__ == '__main__':
    # 检验batch数量很大时， 返回的数量是否相同
    embedding_model = ZhipuEmbeddings()
    r = embedding_model.embed_documents([str(i) for i in range(1000)])
    print(len(r))
